import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import TimeSeriesTransformerModel, TimeSeriesTransformerConfig
import joblib
import os
from datetime import datetime, timedelta


class TimeSeriesDataset(Dataset):
    def __init__(self, X, y, context_length=90, prediction_length=14):
        self.X = X
        self.y = y
        self.context_length = context_length
        self.prediction_length = prediction_length

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        x = self.X[idx]
        y = self.y[idx]

        # 构造 dummy static features (可选)
        static_feat = torch.tensor([0.0], dtype=torch.float32)

        past_time_features = torch.arange(self.context_length).float().unsqueeze(-1)
        future_time_features = torch.arange(self.prediction_length).float().unsqueeze(-1)

        data = {
            "past_values": torch.tensor(x, dtype=torch.float32),
            "past_time_features": past_time_features,
            "future_time_features": future_time_features,
            "static_real_features": static_feat,
            "future_values": torch.tensor(y, dtype=torch.float32).unsqueeze(-1),
        }

        return data


class SalesPredictor:
    def __init__(self, model_path='sales_transformer', context_length=90, prediction_length=14):
        self.model_path = model_path
        self.context_length = context_length
        self.prediction_length = prediction_length
        self.model = None
        self.cutoff_date = None
        self.max_days = context_length  # 商品生命周期90天
        self.forecast_days = prediction_length  # 预测未来14天

        if not os.path.exists(model_path):
            os.makedirs(model_path)

        self.config = TimeSeriesTransformerConfig(
            prediction_length=prediction_length,
            context_length=context_length,
            input_size=1,
            hidden_size=64,
            num_hidden_layers=3,
            num_attention_heads=4,
            lags_sequence=[1, 7, 30],
        )
        self.model = TimeSeriesTransformerModel(self.config)

    def prepare_data(self, raw_df, cutoff_date=None):
        full_data = []
        grouped = raw_df.groupby('spu_code')
        mature_spus = []

        for spu, group in grouped:
            launch_date = group['new_date'].min()
            days_since_launch = (group['date_id'].max() - launch_date).days + 1

            if days_since_launch < self.max_days:
                continue

            mature_spus.append(spu)

            date_range = pd.date_range(start=launch_date, periods=self.max_days, freq='D')
            full_dates = pd.DataFrame({'date_id': date_range, 'spu_code': spu})
            merged = pd.merge(full_dates, group, on=['date_id', 'spu_code'], how='left').fillna(0)
            full_data.append(merged)

        if not mature_spus:
            return np.array([]), np.array([])

        full_df = pd.concat(full_data, ignore_index=True)
        grouped = full_df.groupby('spu_code')

        X, y = [], []
        for spu, group in grouped:
            sales = group.sort_values('date_id')['daily_sale_qty'].values
            if len(sales) < self.max_days:
                sales = np.pad(sales, (0, self.max_days - len(sales)), 'edge')

            for i in range(1, self.max_days - self.forecast_days + 1):
                x_seq = sales[:i]
                y_seq = sales[i:i + self.forecast_days]

                padded_x = np.zeros(self.context_length)
                padded_x[:len(x_seq)] = x_seq

                padded_y = np.zeros(self.prediction_length)
                padded_y[:len(y_seq)] = y_seq

                X.append(padded_x)
                y.append(padded_y)

        return np.array(X), np.array(y)

    def full_train(self, data_path, cutoff_date=None):
        X, y = self.prepare_data(pd.read_excel(data_path))
        if len(X) == 0:
            print("No training data available.")
            return

        dataset = TimeSeriesDataset(X, y, self.context_length, self.forecast_days)
        loader = DataLoader(dataset, batch_size=8, shuffle=True)

        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)

        for epoch in range(50):
            for batch in loader:
                outputs = self.model(
                    past_values=batch["past_values"],
                    past_time_features=batch["past_time_features"],
                    future_time_features=batch["future_time_features"],
                    static_real_features=batch["static_real_features"],
                    future_values=batch["future_values"]
                )
                loss = outputs.loss
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

            print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

        self.model.save_pretrained(self.model_path)
        print(f"Model saved to {self.model_path}")

    def incremental_train(self, new_data_path):
        try:
            self.model = TimeSeriesTransformerModel.from_pretrained(self.model_path)
        except Exception as e:
            print("No existing model found, starting from scratch:", e)

        X, y = self.prepare_data(pd.read_excel(new_data_path))
        if len(X) == 0:
            print("No new training data available.")
            return

        dataset = TimeSeriesDataset(X, y, self.context_length, self.forecast_days)
        loader = DataLoader(dataset, batch_size=8, shuffle=True)

        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-5)

        for epoch in range(10):
            for batch in loader:
                outputs = self.model(
                    past_values=batch["past_values"],
                    past_time_features=batch["past_time_features"],
                    future_time_features=batch["future_time_features"],
                    static_real_features=batch["static_real_features"],
                    future_values=batch["future_values"]
                )
                loss = outputs.loss
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

            print(f"Incremental Epoch {epoch+1}, Loss: {loss.item():.4f}")

        self.model.save_pretrained(self.model_path)
        print(f"Model incrementally trained and saved to {self.model_path}")

    def predict(self, spu_data):
        if self.model is None:
            try:
                self.model = TimeSeriesTransformerModel.from_pretrained(self.model_path)
            except Exception as e:
                raise ValueError("Model not trained or loaded.") from e

        spu_data['date_id'] = pd.to_datetime(spu_data['date_id'], format='%Y-%m-%d')
        spu_data['new_date'] = pd.to_datetime(spu_data['new_date'], format='%Y-%m-%d')

        launch_date = spu_data['new_date'].min()
        last_observed_date = spu_data['date_id'].max()
        current_day = (last_observed_date - launch_date).days

        if current_day >= self.max_days - self.forecast_days:
            print("Exceeds lifecycle.")
            return 0

        full_dates = pd.date_range(start=launch_date, periods=self.max_days, freq='D')
        full_days_df = pd.DataFrame({
            'date_id': full_dates,
            'days_since_launch': range(0, self.max_days)
        })

        spu_data['days_since_launch'] = (spu_data['date_id'] - launch_date).dt.days
        merged = pd.merge(full_days_df, spu_data, on=['date_id', 'days_since_launch'], how='left').fillna(0)
        sales = merged['daily_sale_qty'].values[:self.context_length]

        with torch.no_grad():
            past_values = torch.tensor(sales, dtype=torch.float32).unsqueeze(0)
            past_time_features = torch.arange(self.context_length).float().unsqueeze(0).unsqueeze(-1)
            future_time_features = torch.arange(self.forecast_days).float().unsqueeze(0).unsqueeze(-1)
            static_feat = torch.tensor([0.0])

            outputs = self.model.generate(
                past_values=past_values,
                past_time_features=past_time_features,
                future_time_features=future_time_features,
                static_real_features=static_feat,
                num_return_sequences=1
            )

            forecast = outputs.sequences.mean().item()
            return max(0, forecast)


# 示例用法
if __name__ == "__main__":
    predictor = SalesPredictor()

    # 全量训练
    predictor.full_train('sales.xlsx')

    # 预测示例
    sales_array = [1823, 258, 115, 92, 87, 76, 63, 53, 66, 49, 58, 52, 93, 69, 62, 53, 43, 34, 64, 66, 71, 48, 58, 58, 39,
                   42, 69, 63, 53, 62]
    for i in range(len(sales_array)):
        sub_sales_array = sales_array[:i + 1]
        known_days = len(sub_sales_array)
        new_spu_data = pd.DataFrame({
            'spu_code': ['LO2412090'] * known_days,
            'date_id': pd.date_range(start='2025-01-10', periods=known_days, freq='D'),
            'new_date': ['2025-01-10'] * known_days,
            'daily_sale_qty': sub_sales_array
        })

        forecast = predictor.predict(new_spu_data)
        print(f"上新后第{i+1:02}天销量预测: {forecast:.2f}")
